{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1be44673-1029-4157-86d2-c66587634377",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "sys.path.append(os.path.abspath(os.path.pardir))\n",
    "from argparse import ArgumentParser\n",
    "\n",
    "from voltaml.compile import VoltaNLPCompile\n",
    "from voltaml.inference import nlp_performance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c28fe637-35c3-4211-bb47-ef4c721e4e1c",
   "metadata": {},
   "outputs": [],
   "source": [
    "device=None #choices=[\"cpu\", \"cuda\"]\n",
    "backend=[\"tensorrt\",\"onnx\"] #choices=[\"onnx\",\"tensorrt\"]\n",
    "seq_len=[1, 1, 1] # \"sequence lengths to optimize for (min, optimal, max). Used by TensorRT and benchmarks.\"\n",
    "auth_token=None\n",
    "output='models'\n",
    "task=\"classification\"# [\"classification\", \"embedding\", \"text-generation\", \"token-classification\", \"question-answering\"]\n",
    "batch_size=[1,1,1] #\"batch sizes to optimize for (min, optimal, max). Used by TensorRT and benchmarks.\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b73b45a-ce3e-46aa-b07f-29ea22a99c99",
   "metadata": {},
   "outputs": [],
   "source": [
    "model='bert-base-uncased'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5ef698c-80f1-4ac8-ab3a-168feb211685",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "VoltaNLPCompile(device='cuda',\n",
    "               backend=backend,\n",
    "               seq_len=seq_len,\n",
    "               model=model,\n",
    "               task=task,\n",
    "               )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "287072a4-2241-4a58-a06e-977c4538ebd0",
   "metadata": {},
   "outputs": [],
   "source": [
    "timings = nlp_performance(device='cuda',\n",
    "               backend=backend,\n",
    "               seq_len=seq_len,\n",
    "               model=model,\n",
    "               task=task)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e18b1cc-8f9f-44c1-8af8-c71b1a01fcf6",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d48d917b-3754-4413-9d51-ce058ed90c00",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
